#include "symbol.h"

#include <iostream>

#include "semantics.h"
#include "utils.h"

#define INFO COLOR(5)
#define NAME BLUE
#define ARGUMENT GREY << "| " << DEF
#define VALUE GREEN

ExtFunction::ExtFunction(const std::string& name, char ret,
                         const std::string& arg) {
    this->name = name;
    auto getType = [&](char c) -> Type {
        bool isPtr = std::isupper(c);
        c = std::toupper(c);
        if (c == 'I') return {Type::Basic::Int, isPtr};
        if (c == 'F') return {Type::Basic::Float, isPtr};
        return {Type::Basic::Void, isPtr};
    };
    retType = getType(ret).base;
    if (arg == "...") {
        hasVA = true;
    } else {
        hasVA = false;
        for (auto c : arg) argType.push_back(getType(c));
    }
}

Symbol* Symbol::CreateConstant(const std::string& name, ArithValue value) {
    auto* ret = new Symbol();
    ret->name = name;
    ret->allocMem = 0;
    auto* info = new Semantics::Constant(value);
    ret->info = info;
    return ret;
}

Symbol* Symbol::CreateArray(Type::Basic type, const std::string& name,
                            const std::vector<int>& shape, int length) {
    auto* ret = new Symbol();
    ret->name = name;
    ret->allocMem = length;
    for (auto x : shape) ret->allocMem *= x;
    auto* info = new Semantics::Pointer(type, shape);
    info->symbol = ret;
    ret->info = info;
    return ret;
}

Symbol* Symbol::CreateLValue(Type::Basic type, const std::string& name) {
    auto* ret = new Symbol();
    ret->name = name;
    ret->allocMem = 1;
    auto* info = new Semantics::LValue(type);
    info->symbol = ret;
    ret->info = info;
    return ret;
}

Symbol* Symbol::CreateFunction(Type::Basic type, const std::string& name,
                               AST::ASTNode* func) {
    auto* ret = new Symbol();
    ret->name = name;
    ret->allocMem = 0;
    auto* info = new Semantics::Function(type);
    info->func = func;
    ret->info = info;
    return ret;
}

void SymbolMap::Add(Symbol* sym) {
    std::string name = sym->name;
    if (!syms[name].empty() && syms[name].back().first == depth) {
        Utils::Err("redeclaration of symbol '%s'", name.c_str());
        return;
    }
    syms[name].emplace_back(depth, sym);
    sym->name = name;
    if (GetContext() == Context::Global) {
        // ignore const int
        sym->isGlobal = true;
        if (!sym->info->As<Semantics::Constant>()) {
            Utils::Assert(!global.count(name), "");
            order.push_back(name);
            global[name] = sym;
        }
    }
}

Symbol* SymbolMap::Find(const std::string& name) {
    if (syms.count(name) && !syms[name].empty())
        return syms[name].back().second;
    if (global.count(name)) return global[name];
    if (ext.count(name)) return ext[name];
    Utils::Err("unrecognized symbol '%s'", name.c_str());
    return nullptr;
}

void SymbolMap::CreateEnv(Context c) {
    ++depth;
    context.push_back(c);
}

void SymbolMap::ExitEnv() {
    --depth;
    context.pop_back();
    for (auto& [name, syms] : syms) {
        while (!syms.empty()) {
            auto [d, sym] = syms.back();
            if (d <= depth) break;
            syms.pop_back();
        }
    }
}

void SymbolMap::PrintGlobal() {
    for (auto name : order) {
        auto sym = global[name];
        std::cout << DEF "\n" << name << ": " << sym->allocMem << ", ";
        sym->info->Print(0);
        if (auto info = sym->info->As<Semantics::LValue>()) {
            std::cout << VALUE << info->init << "\n";
        } else if (auto info = sym->info->As<Semantics::Pointer>()) {
            for (auto [x, v] : *info->init) {
                std::cout << NAME "Offset " << x << ": ";
                std::cout << VALUE << v << "\n";
            }
        } else if (auto info = sym->info->As<Semantics::Function>()) {
            info->func->Print(0);
        }
        std::cout << "\n";
    }
}
